import pandas as pd

# 绘图
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib import rcParams

# 设置matplotlib后端为Agg（非交互式），避免GUI相关错误
plt.switch_backend('Agg')
import tushare as ts
import ts.common.base_profile as pf

from scraper.fetch_em_rzrq_total import fetch_one_page as rzrq_fetch_one_page


def init_tushare():
    """初始化Tushare接口"""
    ts.set_token(pf.token)
    return ts.pro_api()


def fetch_index_data(pro, ts_code, start_date="2010101", end_date="20250806"):
    """获取指数数据"""
    return pro.index_daily(**{
        "ts_code": ts_code,
        "start_date": start_date,
        "end_date": end_date,
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code", "trade_date", "close", "open", "high", "low",
        "pre_close", "change", "pct_chg", "vol", "amount"
    ])


def process_trading_amount(df_ssz, df_sesz):
    """处理交易额数据"""
    # 合并df_ssz和df_sesz，按trade_date连接，分别重命名amount列
    df_ssz_ = df_ssz[["trade_date", "amount"]].rename(columns={"amount": "ssz_amount"})
    df_sesz_ = df_sesz[["trade_date", "amount"]].rename(columns={"amount": "sesz_amount"})
    df_merge = pd.merge(df_ssz_, df_sesz_, on="trade_date", how="outer")

    df_merge["amount_sum"] = df_merge[["ssz_amount", "sesz_amount"]].fillna(0).sum(axis=1).round(2)
    # 将千元转换为元（乘以1000）
    df_merge["amount_sum"] = df_merge["amount_sum"] * 1000
    df_merge = df_merge[["trade_date", "ssz_amount", "sesz_amount", "amount_sum"]]

    return df_merge


def process_rzrq_data():
    """处理融资融券数据"""
    df_rzrq_total = pd.read_csv("../result/rzrq_total.csv")
    df_rzrq_increment = rzrq_fetch_one_page(1)  # 增量数据

    # 确保df_rzrq_increment是DataFrame类型
    if isinstance(df_rzrq_increment, list):
        df_rzrq_increment = pd.DataFrame(df_rzrq_increment)

    # 合并增量到全量，按交易日期去重，若重复以增量为准
    df_rzrq_total = pd.concat([df_rzrq_total, df_rzrq_increment], ignore_index=True)

    # 确保交易日期列是字符串格式
    df_rzrq_total['交易日期'] = df_rzrq_total['交易日期'].astype(str)

    df_rzrq_total = df_rzrq_total.sort_values('交易日期').drop_duplicates(subset=['交易日期'], keep='last')
    # 将"交易日期"列格式化为"YYYYMMDD"格式，使用更灵活的日期解析
    df_rzrq_total['交易日期'] = pd.to_datetime(df_rzrq_total['交易日期'], format='mixed', errors='coerce').dt.strftime('%Y%m%d')

    # 保存合并后的结果
    df_rzrq_total.to_csv("../result/rzrq_total.csv", index=False)

    return df_rzrq_total


def merge_and_analyze_data(df_merge, df_rzrq_total):
    """合并数据并进行分析"""
    # 将df_merge和df_rzrq_total按trade_date和交易日期合并
    df_rzrq_total_ = df_rzrq_total.rename(columns={"交易日期": "trade_date"})
    df_merged_all = pd.merge(df_merge, df_rzrq_total_, on="trade_date", how="outer")

    # 新增一列"融资买入额/amount_sum"
    df_merged_all["融资买入额/交易额"] = df_merged_all["融资买入额"] / df_merged_all["amount_sum"]
    # 转换为百分数
    df_merged_all["融资买入额/交易额"] = df_merged_all["融资买入额/交易额"] * 100

    # 按trade_date逆序排序并保存
    df_merged_all = df_merged_all.sort_values("trade_date", ascending=False)
    df_merged_all.to_csv("../result/rzrq_analysis.csv", index=False)

    return df_merged_all


def create_analysis_chart(df_merged_all):
    """创建分析图表"""
    # 设置中文字体
    rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
    rcParams['axes.unicode_minus'] = False

    # 创建图形和双轴
    fig, ax1 = plt.subplots(figsize=(15, 8))
    ax2 = ax1.twinx()

    # 转换日期格式
    df_plot = df_merged_all.copy()
    df_plot['trade_date'] = pd.to_datetime(df_plot['trade_date'])

    # 绘制左轴数据（融资买入额/交易额 和 融资余额余额占流通市值比）
    line1 = ax1.plot(df_plot['trade_date'], df_plot['融资买入额/交易额'],
                     color='blue', linewidth=2, label='融资买入额/交易额(%)')
    line2 = ax1.plot(df_plot['trade_date'], df_plot['融资余额余额占流通市值比'],
                     color='red', linewidth=2, label='融资余额占流通市值比(%)')

    # 绘制右轴数据（沪深300收盘）
    line3 = ax2.plot(df_plot['trade_date'], df_plot['沪深300收盘'],
                     color='green', linewidth=2, label='沪深300收盘')

    # 设置标签和标题
    ax1.set_xlabel('日期')
    ax1.set_ylabel('百分比 (%)', color='black')
    ax2.set_ylabel('沪深300收盘', color='green')

    # 设置标题
    plt.title('融资融券数据分析', fontsize=16, fontweight='bold')

    # 设置图例
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')

    # 格式化x轴日期
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
    ax1.xaxis.set_major_locator(mdates.MonthLocator(interval=3))

    # 旋转x轴标签
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)

    # 设置网格
    ax1.grid(True, alpha=0.3)

    # 调整布局
    plt.tight_layout()

    # 保存图片
    plt.savefig('../result/rzrq_analysis_chart.png', dpi=300, bbox_inches='tight')
    plt.close()  # 关闭图形以释放内存


def main():
    """主函数"""
    # 初始化Tushare
    pro = init_tushare()

    # 获取指数数据
    df_ssz = fetch_index_data(pro, "000001.SH")
    df_sesz = fetch_index_data(pro, "399001.SZ")

    # 处理交易额数据
    df_merge = process_trading_amount(df_ssz, df_sesz)
    df_merge.to_csv("../result/trade_amount.csv", index=False)

    # 处理融资融券数据
    df_rzrq_total = process_rzrq_data()

    # 合并数据并分析
    df_merged_all = merge_and_analyze_data(df_merge, df_rzrq_total)

    # 创建分析图表
    create_analysis_chart(df_merged_all)

    print("融资融券数据分析完成！")


if __name__ == "__main__":
    main()
